#include "solver.h"

void tile_dgemv(const TILE *T, const double *x, const double alpha, double *y)
{
    if (!T || !x || !y)
        return;
    cblas_dgemv(CblasRowMajor, CblasNoTrans, T->m, T->n, alpha, T->a, T->n, x, 1, 1.0, y, 1);
}

void tile_dtrsvLNU(const TILE *T, double *b)
{
    if (!T || !b)
        return;
    cblas_dtrsv(CblasRowMajor, CblasLower, CblasNoTrans, CblasUnit, T->n, T->a, T->n, b, 1);
}

void tile_dtrsvUNN(const TILE *T, double *b)
{
    if (!T || !b)
        return;
    cblas_dtrsv(CblasRowMajor, CblasUpper, CblasNoTrans, CblasNonUnit, T->n, T->a, T->n, b, 1);
}

void block_dgemv(const BLOCK *B, const double *x, const double alpha, double *y, const int inc)
{
    if (!B || !x || !y)
        return;

    for (int i = 0; i < B->m; i++)
    {
        for (int j = 0; j < B->n; j++)
        {
            TILE *T = B->tile[i * B->n + j];
            if (!T)
                continue;
            tile_dgemv(T, &x[j * inc], alpha, &y[i * inc]);
        }
    }
}

void block_dtrsvLNU(const BLOCK *B, double *b, const int incb)
{
    if (!B || !b)
        return;

    TILE *T = NULL;
    for (int i = 0; i < B->m; i++)
    {
        for (int j = 0; j < i; j++)
        {
            T = B->tile[i * B->n + j];
            if (!T)
                continue;
            tile_dgemv(T, &b[j * incb], -1.0, &b[i * incb]);
        }
        T = B->tile[i * B->n + i];
        if (!T)
            continue;
        tile_dtrsvLNU(T, &b[i * incb]);
    }
}

void block_dtrsvUNN(const BLOCK *B, double *b, int incb)
{
    if (!B || !b)
        return;

    TILE *T = NULL;
    for (int j = B->n - 1; j >= 0; j--)
    {
        T = B->tile[j * B->n + j];
        if (!T)
            continue;
        tile_dtrsvUNN(T, &b[j * incb]);
        for (int i = 0; i < j; i++)
        {
            T = B->tile[i * B->n + j];
            if (!T)
                continue;
            tile_dgemv(T, &b[j * incb], -1.0, &b[i * incb]);
        }
    }
}

void Lsolve(const MATRIX *A, double *b)
{
    if (!A || !b)
        return;

    BLOCK *B = NULL;
    int incb = A->block_size * A->tile_size;
    for (int i = 0; i < A->m; i++)
    {
        for (int j = 0; j < i; j++)
        {
            B = A->block[i * A->n + j];
            if (!B)
                continue;
            block_dgemv(B, &b[j * incb], -1.0, &b[i * incb], A->tile_size);
        }
        B = A->block[i * A->n + i];
        if (!B)
            continue; // 必须，否则b会因为越界段错误
        block_dtrsvLNU(B, &b[i * incb], A->tile_size);
    }
}

void Usolve(const MATRIX *A, double *b)
{
    if (!A || !b)
        return;

    BLOCK *B = NULL;
    int incb = A->block_size * A->tile_size;
    for (int j = A->n - 1; j >= 0; j--)
    {
        B = A->block[j * A->n + j];
        if (!B)
            continue;
        block_dtrsvUNN(B, &b[j * incb], A->tile_size);
        for (int i = 0; i < j; i++)
        {
            B = A->block[i * A->n + j];
            if (!B)
                continue;
            block_dgemv(B, &b[j * incb], -1.0, &b[i * incb], A->tile_size);
        }
    }
}

void solve(MATRIX *A, double *b)
{
    if (!A || !b)
        return;
    Lsolve(A, b);
    Usolve(A, b);
}

void gemv(MATRIX *A, const double *x, double *y)
{
    if (!A || !x || !y)
        return;

    int inc = A->block_size * A->tile_size;
    for (int i = 0; i < A->m; i++)
    {
        for (int j = 0; j < A->n; j++)
        {
            BLOCK *B = A->block[i * A->n + j];
            if (!B)
                continue;
            block_dgemv(B, &x[j * inc], 1.0, &y[i * inc], A->tile_size);
        }
    }
}
